{
  "cells": [
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "rSwOkL0XIKWa"
      },
      "outputs": [],
      "source": [
        "!pip install -q git+https://github.com/google/tunix\n",
        "!pip install -q git+https://github.com/google/qwix\n",
        "!pip uninstall -q -y flax\n",
        "!pip install --no-cache-dir git+https://github.com/google/flax.git\n",
        "!pip install -q huggingface_hub\n",
        "!pip install -q humanize"
      ]
    },
    {
      "metadata": {
        "id": "Crz-9Yt7IO0v"
      },
      "cell_type": "code",
      "source": [
        "!huggingface-cli login"
      ],
      "outputs": [],
      "execution_count": null
    },
    {
      "metadata": {
        "id": "QJjzmkSCIRqv"
      },
      "cell_type": "code",
      "source": [
        "import functools\n",
        "from flax import nnx\n",
        "from huggingface_hub import snapshot_download\n",
        "import humanize\n",
        "import jax\n",
        "from tunix.models.gemma import model as model_lib\n",
        "from tunix.models.gemma import params_safetensors as params_lib"
      ],
      "outputs": [],
      "execution_count": null
    },
    {
      "metadata": {
        "id": "cr7lQ0PmU1Fd"
      },
      "cell_type": "code",
      "source": [
        "def show_hbm_usage():\n",
        "  \"\"\"Displays memory usage per device.\"\"\"\n",
        "  fmt_size = functools.partial(humanize.naturalsize, binary=True)\n",
        "\n",
        "  print(\"\\n--- TPU HBM Usage ---\")\n",
        "  for i, d in enumerate(jax.local_devices()):\n",
        "    stats = d.memory_stats()\n",
        "    used = stats.get(\"bytes_in_use\", 0)\n",
        "    limit = stats.get(\"bytes_limit\", 0)\n",
        "\n",
        "    hbm_used = stats.get(\"device:0:HBM0:bytes_in_use\", used)\n",
        "    hbm_limit = stats.get(\"device:0:HBM0:bytes_limit\", limit)\n",
        "\n",
        "    # Fallback if specific HBM stats not available\n",
        "    if hbm_limit == 0:\n",
        "      hbm_used = used\n",
        "      hbm_limit = limit\n",
        "\n",
        "    percentage = (hbm_used / hbm_limit * 100) if hbm_limit > 0 else 0\n",
        "\n",
        "    print(\n",
        "        f\"Device {i} ({d.device_kind}): Using {fmt_size(hbm_used)} /\"\n",
        "        f\" {fmt_size(hbm_limit)} ({percentage:.2f}%)\"\n",
        "    )\n",
        "\n",
        "  print(\"--- End HBM Usage ---\")"
      ],
      "outputs": [],
      "execution_count": null
    },
    {
      "metadata": {
        "id": "zcP25ZQ3TsZ7"
      },
      "cell_type": "markdown",
      "source": [
        "# Download the weights and load into TPU"
      ]
    },
    {
      "metadata": {
        "id": "yN5jp8b9VBXN"
      },
      "cell_type": "code",
      "source": [
        "model_id = \"google/gemma-2-2b-it\"\n",
        "ignore_patterns = [\n",
        "    \"*.pth\",  # Ignore PyTorch .pth weight files\n",
        "]\n",
        "print(f\"Downloading {model_id} from Hugging Face...\")\n",
        "local_model_path = snapshot_download(\n",
        "    repo_id=model_id, ignore_patterns=ignore_patterns\n",
        ")\n",
        "print(f\"Model successfully downloaded to: {local_model_path}\")"
      ],
      "outputs": [],
      "execution_count": null
    },
    {
      "metadata": {
        "id": "IJQKqNPaVPz6"
      },
      "cell_type": "code",
      "source": [
        "print(\"\\n--- HBM Usage BEFORE Model Load ---\")\n",
        "show_hbm_usage()"
      ],
      "outputs": [],
      "execution_count": null
    },
    {
      "metadata": {
        "id": "fZ3hUjf4VThk"
      },
      "cell_type": "code",
      "source": [
        "MODEL_CP_PATH = local_model_path\n",
        "\n",
        "config = model_lib.ModelConfig.gemma2_2b()\n",
        "MESH = [(1, 1), (\"fsdp\", \"tp\")]  # update this based on your # TPU devices\n",
        "mesh = jax.make_mesh(*MESH, axis_types=(jax.sharding.AxisType.Auto,) * len(MESH[0]))\n",
        "with mesh:\n",
        "  gemma = params_lib.create_model_from_safe_tensors(MODEL_CP_PATH, config, mesh)\n",
        "  nnx.display(gemma)"
      ],
      "outputs": [],
      "execution_count": null
    },
    {
      "metadata": {
        "id": "tc_a5GUEVeXf"
      },
      "cell_type": "code",
      "source": [
        "print(\"\\n--- HBM Usage AFTER Model Load ---\")\n",
        "show_hbm_usage()"
      ],
      "outputs": [],
      "execution_count": null
    },
    {
      "metadata": {
        "id": "jZCKBEjvTm_A"
      },
      "cell_type": "markdown",
      "source": [
        "# Run inference"
      ]
    },
    {
      "metadata": {
        "id": "31D_52F_TlTs"
      },
      "cell_type": "code",
      "source": [
        "from transformers import AutoTokenizer\n",
        "\n",
        "tokenizer = AutoTokenizer.from_pretrained(MODEL_CP_PATH)"
      ],
      "outputs": [],
      "execution_count": null
    },
    {
      "metadata": {
        "id": "kb--p9NxT7x1"
      },
      "cell_type": "code",
      "source": [
        "from tunix.generate import sampler\n",
        "\n",
        "\n",
        "def templatize(prompts):\n",
        "  out = []\n",
        "  for p in prompts:\n",
        "    out.append(\n",
        "        tokenizer.apply_chat_template(\n",
        "            [\n",
        "                {\"role\": \"user\", \"content\": p},\n",
        "            ],\n",
        "            tokenize=False,\n",
        "            add_generation_prompt=True,\n",
        "            enable_thinking=True,\n",
        "        )\n",
        "    )\n",
        "  return out\n",
        "\n",
        "\n",
        "inputs = templatize([\n",
        "    \"which is larger 9.9 or 9.11?\",\n",
        "    \"如何制作月饼?\",\n",
        "    \"tell me your name, respond in Spanish\",\n",
        "])\n",
        "\n",
        "sampler = sampler.Sampler(\n",
        "    gemma,\n",
        "    tokenizer,\n",
        "    sampler.CacheConfig(\n",
        "        cache_size=256,\n",
        "        num_layers=config.num_layers,\n",
        "        num_kv_heads=config.num_kv_heads,\n",
        "        head_dim=config.head_dim,\n",
        "    ),\n",
        ")\n",
        "out = sampler(inputs, max_generation_steps=128, echo=True)\n",
        "\n",
        "for t in out.text:\n",
        "  print(t)\n",
        "  print(\"*\" * 30)"
      ],
      "outputs": [],
      "execution_count": null
    }
  ],
  "metadata": {
    "colab": {
      "private_outputs": true
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}
